/*
 *   Copyright 2023 <a href="mailto:asialjim@hotmail.com">Asial Jim</a>
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */
package com.asialjim.encommunicate.base.rsa;

import org.apache.commons.io.IOUtils;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Objects;

/**
 * 基于 RSA 的加解密工具
 *
 * @author Copyright © <a href="mailto:asialjim@hotmail.com">Asial Jim</a>   Co., LTD
 * @version 1.0
 * @since 2023/2/24, &nbsp;&nbsp; <em>version:1.0</em>, &nbsp;&nbsp; <em>java version:17</em>
 */
@SuppressWarnings("unused")
public interface RSAEncrypt {

    /**
     * 加密
     *
     * @param source {@link String 待加密内容}
     * @return {@link String 加密结果}
     * @since 2023/2/24
     */
    String encrypt(String source);

    /**
     * 加密
     *
     * @param source {@link InputStream 待加密内容}
     * @return {@link java.lang.reflect.Array 加密结果}
     * @since 2023/2/24
     */
    byte[] encrypt(InputStream source);

    /**
     * 加密
     *
     * @param source {@link java.lang.reflect.Array 待加密内容}
     * @return {@link java.lang.reflect.Array 加密结果}
     * @since 2023/2/24
     */
    byte[] encrypt(byte[] source);


    default String encrypt(String source, String pubKey) {
        byte[] encrypt = encrypt(source.getBytes(StandardCharsets.UTF_8), pubKey);
        return Base64.getEncoder().encodeToString(encrypt);
    }

    default byte[] encrypt(InputStream inputStream, String pubKey) {
        ByteArrayOutputStream bos = null;
        try {
            byte[] pubKeyBytes = Base64.getDecoder().decode(pubKey);
            RSAPublicKey publicKey = (RSAPublicKey) KeyFactory.getInstance(RSAKeyPairProperty.keyAlgorithm()).generatePublic(new X509EncodedKeySpec(pubKeyBytes));
            int maxSourceLength = (publicKey.getModulus().bitLength() >> 3) - 11;
            if (inputStream.available() <= maxSourceLength)
                return encrypt(inputStream.readNBytes(inputStream.available()), pubKey);

            Cipher cipher = Cipher.getInstance(RSAKeyPairProperty.keyAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);

            bos = new ByteArrayOutputStream();
            byte[] block;
            do {
                block = inputStream.readNBytes(maxSourceLength);
                if (Objects.isNull(block) || block.length == 0) break;

                bos.write(cipher.doFinal(block));
            } while (block.length >= maxSourceLength);

            return bos.toByteArray();
        } catch (IOException | InvalidKeySpecException | NoSuchAlgorithmException |
                 NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {

            throw new RuntimeException(e);
        } finally {
            IOUtils.closeQuietly(bos);
        }
    }

    default byte[] encrypt(byte[] source, String pubKey) {
        ByteArrayOutputStream bos = null;
        ByteArrayInputStream bis = null;
        try {
            byte[] pubKeyBytes = Base64.getDecoder().decode(pubKey);
            RSAPublicKey publicKey = (RSAPublicKey) KeyFactory.getInstance(RSAKeyPairProperty.keyAlgorithm()).generatePublic(new X509EncodedKeySpec(pubKeyBytes));
            Cipher cipher = Cipher.getInstance(RSAKeyPairProperty.keyAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);

            int maxSourceLength = (publicKey.getModulus().bitLength() >> 3) - 11;
            if (source.length <= maxSourceLength)
                return cipher.doFinal(source);

            bos = new ByteArrayOutputStream();
            bis = new ByteArrayInputStream(source);

            byte[] block;
            do {
                block = bis.readNBytes(maxSourceLength);
                if (Objects.isNull(block) || block.length == 0) break;

                bos.write(cipher.doFinal(block));
            } while (block.length >= maxSourceLength);

            return bos.toByteArray();
        } catch (NoSuchPaddingException | IllegalBlockSizeException | InvalidKeySpecException |
                 NoSuchAlgorithmException | BadPaddingException | InvalidKeyException | IOException e) {

            throw new RuntimeException(e);
        } finally {
            IOUtils.closeQuietly(bos);
            IOUtils.closeQuietly(bis);
        }
    }
}